/*
    Ask the user to run the program with admin privileged.
    
    Author @5mukx
 */

use std::process::Command;
use std::ptr::null_mut;
use winapi::um::shellapi::ShellExecuteW;
use winapi::um::winnt::{HANDLE, TOKEN_ELEVATION, TOKEN_QUERY};
use winapi::um::processthreadsapi::OpenProcessToken;
use winapi::um::securitybaseapi::GetTokenInformation;
use winapi::um::winuser::{MessageBoxW, MB_ICONERROR, MB_OK};

fn is_running_as_admin() -> bool {
    unsafe {
        let mut token_handle: HANDLE = null_mut();
        if OpenProcessToken(
            winapi::um::processthreadsapi::GetCurrentProcess(),
            TOKEN_QUERY,
            &mut token_handle,
        ) == 0
        {
            return false;
        }

        let mut elevation: TOKEN_ELEVATION = std::mem::zeroed();
        let mut size = 0;

        let success = GetTokenInformation(
            token_handle,
            winapi::um::winnt::TokenElevation,
            &mut elevation as *mut _ as *mut _,
            std::mem::size_of::<TOKEN_ELEVATION>() as u32,
            &mut size,
        ) != 0;

        winapi::um::handleapi::CloseHandle(token_handle);

        success && elevation.TokenIsElevated != 0
    }
}

fn run_as_admin() {
    unsafe {
        let current_exe = std::env::current_exe().unwrap();
        let current_exe_wide: Vec<u16> = current_exe
            .to_string_lossy()
            .encode_utf16()
            .chain(Some(0))
            .collect();

        ShellExecuteW(
            null_mut(),
            "runas\0".encode_utf16().chain(Some(0)).collect::<Vec<u16>>().as_ptr(),
            current_exe_wide.as_ptr(),
            null_mut(),
            null_mut(),
            winapi::um::winuser::SW_SHOWNORMAL,
        );
    }
}

fn main() {
    // Avoid infinite relaunches
    let elevated_env_var = "RUNNING_ELEVATED";

    if std::env::var(elevated_env_var).is_err() {
        if !is_running_as_admin() {
            println!("This program requires administrator privileges.");
            unsafe {
                MessageBoxW(
                    null_mut(),
                    "Please run the program as Administrator.".encode_utf16().chain(Some(0)).collect::<Vec<u16>>().as_ptr(),
                    "Administrator Access Required".encode_utf16().chain(Some(0)).collect::<Vec<u16>>().as_ptr(),
                    MB_ICONERROR | MB_OK,
                );
            }

            std::env::set_var(elevated_env_var, "1");
            run_as_admin();

            std::process::exit(0);
        }
    }

    // Your privileged operations here
    println!("Running with administrator privileges!");
    
}

